# -*- coding: utf-8 -*-
"""
Created on 2025/7/5 16:11
---------
@summary: 拆分且打乱isbn
---------
@author: rui_kai
"""
import csv
import math
import os
import queue
import random
import threading
from concurrent.futures import ThreadPoolExecutor

from rich.console import Console

from utils.excel_handler import ExcelHandlerNew
from utils.excel_read import ExcelRead
from utils.tools import set_save_path


class SplitIsbnAndShuffle:
    counter_lock = threading.Lock()

    def __init__(self, file_path):
        self.file_counter = 0
        self.file_path = file_path
        self.console = Console()
        self.q = queue.Queue()

        self.one_file_max_row = None
        self.split_num = None
        self.data_type = None
        self.data_list = []
        self.write_excel_path = os.path.splitext(os.path.basename(file_path))[0]
        self.run()

    def __start_callback(self):
        while True:
            spilt_num = self.console.input(f"[bold green]请输入拆分文件数量:[/bold green]")
            try:
                self.split_num = int(spilt_num)
            except ValueError:
                self.console.print("[bold red]请输入数字[/bold red]")
                continue
            break

        self.console.print(f"[bold green]任务已经开始..请耐心等待...[/bold green]")

    def __set_q(self):
        for i in self.data_list:
            self.q.put(i)
        self.push_complete = True

    def __get_data_type(self):
        if self.file_path.endswith('.txt'):
            self.data_type = "txt"
        elif self.file_path.endswith('.xlsx'):
            self.data_type = "xlsx"
        elif self.file_path.endswith('.csv'):
            self.data_type = "csv"
        else:
            self.console.print("[bold red]只支持txt、xlsx、csv格式的文件[/bold red]")
            exit()

    def _read_txt(self):
        with open(self.file_path, "r", encoding="utf-8") as f:
            data = f.read()
        self.data_list = [i for i in data.split("\n") if i]

    def __read_xlsx(self):
        excel = ExcelRead(self.file_path)
        excel.load_workbook()
        for row in excel.sheet.iter_rows(min_col=1, max_col=1):
            for cell in row:
                self.data_list.append(cell.value)

    def __read_csv(self):
        with open(self.file_path, "r", encoding="utf-8", errors="replace") as f:
            data = csv.reader(f)
            self.data_list = [i[0] for i in data]

    def thread_start(self):
        random.shuffle(self.data_list)
        if self.data_type == "txt":
            self._read_txt()
        elif self.data_type == "xlsx":
            self.__read_xlsx()
        elif self.data_type == "csv":
            self.__read_csv()

        # 计算每份数据量
        thread_count = min(5, self.split_num)
        self.one_file_max_row = min(math.ceil(len(self.data_list) / self.split_num), 500000)

        self.console.print(f"本次任务有 {len(self.data_list)} 条数据 共存 {self.split_num} 份")
        if len(self.data_list) == 0:
            self.console.print("[bold red]没有数据可处理[/bold red]")
            return

        threading.Thread(target=self.__set_q, daemon=True).start()
        self.thread_export(thread_count)

    def thread_export(self, thread_count):
        with ThreadPoolExecutor(max_workers=thread_count) as executor:
            futures = []
            for i in range(thread_count):
                if self.data_type == "txt":
                    futures.append(executor.submit(self._export_txt))
                else:
                    futures.append(executor.submit(self._export_excel))

            for future in futures:
                future.result()

        self.file_counter = 0

    def _export_txt(self):
        while self.file_counter < self.split_num:
            with self.__class__.counter_lock:  # 线程安全地获取并增加计数器
                current_file_index = self.file_counter
                self.file_counter += 1

            file_name = f"{self.write_excel_path}_{current_file_index}"
            save_file_path = set_save_path(file_name, "txt")
            self.console.print(f"开始写入文本文件：{save_file_path}")
            count = 0
            with open(save_file_path, "w", encoding="utf-8") as f:
                while True:
                    if count >= self.one_file_max_row:
                        break
                    try:
                        item = self.q.get(block=True, timeout=2)
                    except queue.Empty:
                        if self.push_complete and self.q.empty():
                            break
                        continue
                    f.write(f"{item}\n")
                    count += 1

            self.console.print(f"[blue]文本： [yellow]{save_file_path}[/yellow] 写入完成。[/blue]")

    def _export_excel(self):
        while self.file_counter < self.split_num:
            with self.__class__.counter_lock:  # 线程安全地获取并增加计数器
                current_file_index = self.file_counter
                self.file_counter += 1

            file_name = f"{self.write_excel_path}_{current_file_index}"
            save_file_path = set_save_path(file_name, self.data_type)
            self.console.print(f"开始写入表格：{save_file_path}")
            excel = ExcelHandlerNew(save_file_path)

            while True:
                if excel.rows >= self.one_file_max_row:
                    excel.save()
                    break

                try:
                    item = self.q.get(block=True, timeout=2)
                    # print([item])
                    excel.write_row([item])
                    self.q.task_done()
                except queue.Empty:
                    if self.push_complete and self.q.empty():
                        # 数据全部处理完毕
                        if excel.rows > 1:
                            excel.save()
                        break
                    continue

            self.console.print(f"[blue]表格：[yellow]{save_file_path}[/yellow] 写入完成。[/blue]")

    def run(self):
        self.__get_data_type()
        self.__start_callback()
        self.thread_start()

    def set_save_path(self, file_name, suffix):
        pass
